"""
独立的扩散过程可视化脚本
支持在训练外单独运行，可视化：
1. 前向加噪过程（数据 → 噪声）
2. 反向去噪过程（噪声 → 重建）
"""
import argparse
import os
import sys
import numpy as np
import torch

# 添加当前目录到 Python 路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from Squidiff import dist_util
from Squidiff.hic_script_util import (
    get_all_defaults,
    create_model_and_diffusion,
    args_to_dict,
    add_dict_to_argparser,
)
from Squidiff.hic_datasets import create_hic_dataloader
from Squidiff.hic_visualization import (
    visualize_forward_diffusion,
    visualize_reverse_diffusion,
    visualize_comparison,
)


def main():
    parser = argparse.ArgumentParser(description='Hi-C 扩散过程可视化')
    
    # 基础参数
    parser.add_argument('--mode', type=str, default='forward', 
                       choices=['forward', 'reverse', 'both'],
                       help='可视化模式：forward（加噪）, reverse（去噪）, both（两者）')
    parser.add_argument('--model_path', type=str, default='',
                       help='模型检查点路径（reverse 模式必需）')
    parser.add_argument('--data_dir', type=str, required=True,
                       help='Hi-C 数据目录')
    parser.add_argument('--output_dir', type=str, default='./visualizations',
                       help='可视化结果输出目录')
    
    # 数据参数
    parser.add_argument('--timepoints', type=str, default='G1_top400,late_S1_top400',
                       help='时间点（逗号分隔）')
    parser.add_argument('--file_format', type=str, default='npz',
                       choices=['cool', 'txt', 'npy', 'npz'],
                       help='文件格式')
    parser.add_argument('--chromosome', type=str, default='chr1',
                       help='染色体')
    parser.add_argument('--resolution', type=int, default=100000,
                       help='分辨率（bp）')
    parser.add_argument('--window_size', type=int, default=256,
                       help='窗口大小')
    parser.add_argument('--normalization', type=str, default='minmax',
                       help='归一化方法')
    parser.add_argument('--log_transform', type=bool, default=True,
                       help='对数变换')
    
    # 可视化参数
    parser.add_argument('--num_timesteps', type=int, default=16,
                       help='展示的时间步数量')
    parser.add_argument('--num_samples', type=int, default=4,
                       help='可视化样本数')
    parser.add_argument('--timestep_strategy', type=str, default='linspace',
                       choices=['linspace', 'random'],
                       help='时间步选择策略')
    parser.add_argument('--use_mean', action='store_true',
                       help='只显示均值（无随机噪声）')
    parser.add_argument('--log_display', action='store_true', default=True,
                       help='使用 log1p 显示')
    parser.add_argument('--signed_display', action='store_true',
                       help='支持正负值显示（RdBu_r）')
    
    # 模型参数（从默认配置读取）
    defaults = get_all_defaults()
    for k, v in defaults.items():
        if k not in ['data_dir', 'timepoints', 'file_format', 'chromosome', 
                     'resolution', 'window_size', 'normalization', 'log_transform']:
            v_type = type(v)
            if v is None:
                v_type = str
            elif isinstance(v, bool):
                v_type = lambda x: x.lower() in ("yes", "true", "t", "1")
            parser.add_argument(f'--{k}', default=v, type=v_type)
    
    args = parser.parse_args()
    
    print("=" * 80)
    print("Hi-C 扩散过程可视化")
    print("=" * 80)
    print(f"模式: {args.mode}")
    print(f"数据目录: {args.data_dir}")
    print(f"输出目录: {args.output_dir}")
    print(f"时间步数: {args.num_timesteps}")
    print(f"样本数: {args.num_samples}")
    print("=" * 80)
    
    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)
    
    # 加载数据
    print("\n加载数据...")
    dataloader = create_hic_dataloader(
        data_dir=args.data_dir,
        batch_size=args.num_samples,
        timepoints=args.timepoints.split(','),
        file_format=args.file_format,
        chromosome=args.chromosome,
        resolution=args.resolution,
        window_size=args.window_size,
        normalization=args.normalization,
        log_transform=args.log_transform,
        interpolation_prob=0.0,
        augmentation=False,
        num_workers=2,
    )
    
    # 获取一个批次
    batch = next(iter(dataloader))
    print(f"数据批次形状: {batch['hic'].shape}")
    
    # 创建扩散模型
    print("\n创建扩散模型...")
    model_args = {
        'image_size': args.window_size,
        'in_channels': args.in_channels,
        'model_channels': args.model_channels,
        'out_channels': args.out_channels,
        'num_res_blocks': args.num_res_blocks,
        'attention_resolutions': args.attention_resolutions,
        'dropout': args.dropout,
        'channel_mult': args.channel_mult,
        'conv_resample': args.conv_resample,
        'dims': args.dims,
        'num_classes': args.num_classes,
        'use_checkpoint': args.use_checkpoint,
        'use_fp16': args.use_fp16,
        'num_heads': args.num_heads,
        'num_head_channels': args.num_head_channels,
        'num_heads_upsample': args.num_heads_upsample,
        'use_scale_shift_norm': args.use_scale_shift_norm,
        'resblock_updown': args.resblock_updown,
        'use_new_attention_order': args.use_new_attention_order,
        'use_spatial_transformer': args.use_spatial_transformer,
        'transformer_depth': args.transformer_depth,
        'context_dim': args.context_dim,
        'time_cond_dim': args.time_cond_dim,
        'use_transition_anchors': args.use_transition_anchors,
        'cond_drop_prob': args.cond_drop_prob,
        'sem_cond_dim': args.sem_cond_dim,
        'use_semantic_encoder': args.use_semantic_encoder,
        'sem_drop_prob': args.sem_drop_prob,
        'use_vae_encoder': args.use_vae_encoder,
        'vae_encoder_path': args.vae_encoder_path,
        'freeze_vae_encoder': args.freeze_vae_encoder,
    }
    
    diffusion_args = {
        'learn_sigma': args.learn_sigma,
        'diffusion_steps': args.diffusion_steps,
        'noise_schedule': args.noise_schedule,
        'timestep_respacing': args.timestep_respacing,
        'use_kl': args.use_kl,
        'predict_xstart': args.predict_xstart,
        'rescale_timesteps': args.rescale_timesteps,
        'rescale_learned_sigmas': args.rescale_learned_sigmas,
    }
    
    model, diffusion = create_model_and_diffusion(**model_args, **diffusion_args)
    
    # 加载模型权重（如果提供）
    if args.model_path and os.path.exists(args.model_path):
        print(f"\n加载模型权重: {args.model_path}")
        state_dict = torch.load(args.model_path, map_location='cpu')
        model.load_state_dict(state_dict)
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    model.eval()
    
    print(f"\n模型已加载到设备: {device}")
    print(f"扩散步数: {diffusion.num_timesteps}")
    print(f"Beta schedule: {args.noise_schedule}")
    
    # 前向加噪可视化
    if args.mode in ['forward', 'both']:
        print("\n" + "=" * 80)
        print("前向加噪可视化")
        print("=" * 80)
        
        for i in range(min(args.num_samples, batch['hic'].shape[0])):
            x0 = batch['hic'][i, 0].numpy()  # [H, W]
            
            save_path = os.path.join(args.output_dir, f'forward_sample_{i}.png')
            
            cell_id = batch.get('cell_id', ['unknown'])[i]
            timepoint = batch.get('timepoint', ['unknown'])[i]
            suptitle = f"前向扩散过程\nCell: {cell_id} | Timepoint: {timepoint}"
            
            print(f"\n样本 {i}: Cell={cell_id}, Timepoint={timepoint}")
            print(f"保存至: {save_path}")
            
            visualize_forward_diffusion(
                x0=x0,
                betas=diffusion.betas,
                num_timesteps=args.num_timesteps,
                timestep_strategy=args.timestep_strategy,
                use_mean=args.use_mean,
                log_display=args.log_display,
                signed_display=args.signed_display,
                save_path=save_path,
                suptitle=suptitle,
            )
    
    # 反向去噪可视化
    if args.mode in ['reverse', 'both']:
        print("\n" + "=" * 80)
        print("反向去噪可视化")
        print("=" * 80)
        
        if not args.model_path:
            print("警告：未提供模型路径，使用未训练的模型（结果可能不佳）")
        
        shape = (1, 1, args.window_size, args.window_size)
        
        for i in range(min(args.num_samples, batch['hic'].shape[0])):
            bio_time = batch['time'][i:i+1].to(device)
            
            save_path = os.path.join(args.output_dir, f'reverse_sample_{i}.png')
            
            timepoint = batch.get('timepoint', ['unknown'])[i]
            suptitle = f"反向去噪过程\nTimepoint: {timepoint}"
            
            print(f"\n样本 {i}: Timepoint={timepoint}")
            print(f"保存至: {save_path}")
            
            visualize_reverse_diffusion(
                model=model,
                diffusion=diffusion,
                shape=shape,
                device=device,
                num_timesteps=args.num_timesteps,
                timestep_strategy=args.timestep_strategy,
                log_display=args.log_display,
                signed_display=args.signed_display,
                save_path=save_path,
                model_kwargs={'bio_time': bio_time},
                suptitle=suptitle,
            )
    
    print("\n" + "=" * 80)
    print("可视化完成！")
    print(f"结果保存在: {args.output_dir}")
    print("=" * 80)


if __name__ == '__main__':
    main()

